# coding:utf-8
# __author__ = yuan
# __time__ = 2020/4/23
# __file__ = test
# __desc__ =

import os
from scipy import misc
import time
import tensorflow as tf
from tensorflow.python.framework.errors_impl import OutOfRangeError
import numpy as np
import glob
from sklearn.metrics import confusion_matrix
from net import se_resnet
import tensorflow.contrib.slim as slim
from data import _get_data_list,mixup
import pyecharts
import matplotlib.pyplot as plt
os.environ['CUDA_VISIBLE_DEVICES'] = "3"

def plot(mt:np.ndarray,labels):
    plt.matshow(mt)
    plt.colorbar()
    plt.xlabel("预测类型")
    plt.ylabel("真实类型")
    plt.xticks(np.arange(mt.shape[1]),labels)
    plt.yticks(np.arange(mt.shape[1]),labels)
    plt.show()

def test_model():
    img_root = r"/data/soft/javad/COCO/convd/test"
    save_dir = "./models"
    assert os.path.exists(save_dir)
    batch = 3
    H = 900
    W = 900
    numclass = 2
    alpha = 5
    label_map={0:"normal",1:"pneumonia"}
    model_graph = tf.train.latest_checkpoint(save_dir)
    inputs = tf.placeholder(tf.float32, [None, H, W, 3])
    labels = tf.placeholder(tf.uint8, [None, numclass])
    logit = se_resnet(inputs, blocks=[3, 4, 6, 3], numclass=numclass)
    prop = slim.softmax(logit)
    correct_prediction = tf.equal(tf.argmax(prop, 1), tf.argmax(labels, 1))  # predict_op = tf.argmax(y, 1)
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, 'float'))

    saver = tf.train.Saver()
    cpu_c=3
    cpu_config = tf.ConfigProto(intra_op_parallelism_threads=cpu_c,
                                inter_op_parallelism_threads=cpu_c,
                                device_count={'CPU': cpu_c})
    cpu_config.gpu_options.allow_growth = True
    preds= []
    reals = []
    images,labels = _get_data_list(img_root)
    with tf.Session(config=cpu_config) as sess:
        sess.run(tf.global_variables_initializer)
        saver.restore(sess,model_graph)
        for (imp,label) in zip(images,labels):
            img = misc.imread(imp)
            img = np.expand_dims(img)
            pred:np.ndarray = sess.run(prop,feed_dict={inputs:img})
            pred = np.ravel(pred)[0]
            preds.append(pred)
            reals.append(label)
            print(f"{imp} 预测结果:\n"
                  f"预测: {pred}  -- 真实: {label}")

    preds=list(map(lambda x:label_map[x],preds))
    reals=list(map(lambda x:label_map[x],reals))
    confuse_mat=confusion_matrix(reals,preds)

    plot(confuse_mat,list(label_map.values()))
